#豆包生成，将txt形式的标注转换成.json标注
import os
import json
from PIL import Image

def yolo_to_coco(yolo_images_dir, yolo_labels_dir, output_json_path):
    # COCO格式的基本结构
    coco_format = {
        "info": {},
        "licenses": [],
        "categories": [],
        "images": [],
        "annotations": []
    }
    
    # 1. 读取类别（COCO128的类别与COCO一致，共80类）
    # 从labels文件夹中找一个标注文件获取类别数量
    label_files = [f for f in os.listdir(yolo_labels_dir) if f.endswith('.txt')]
    if not label_files:
        raise ValueError("未找到标注文件，请检查labels路径")
    
    # 假设所有标注的类别索引都在0-79之间（COCO标准）
    for i in range(80):
        coco_format["categories"].append({
            "id": i,
            "name": f"class_{i}",  # 实际类别名不影响测试，只要ID正确即可
            "supercategory": "none"
        })
    
    # 2. 处理图像和标注
    image_id = 0
    annotation_id = 0
    
    for img_file in os.listdir(yolo_images_dir):
        if not img_file.endswith(('.jpg', '.jpeg', '.png')):
            continue
        
        # 图像信息
        img_path = os.path.join(yolo_images_dir, img_file)
        img = Image.open(img_path)
        width, height = img.size
        
        coco_format["images"].append({
            "id": image_id,
            "file_name": img_file,
            "width": width,
            "height": height
        })
        
        # 对应标注文件
        label_file = os.path.splitext(img_file)[0] + '.txt'
        label_path = os.path.join(yolo_labels_dir, label_file)
        
        if os.path.exists(label_path):
            with open(label_path, 'r') as f:
                for line in f.readlines():
                    line = line.strip()
                    if not line:
                        continue
                    
                    # YOLO格式：class_id x_center y_center width height（均为归一化值）
                    class_id, x_center, y_center, w, h = map(float, line.split())
                    class_id = int(class_id)
                    
                    # 转换为COCO格式（左上角坐标 + 宽高，绝对像素值）
                    x_min = (x_center - w/2) * width
                    y_min = (y_center - h/2) * height
                    bbox_width = w * width
                    bbox_height = h * height
                    
                    coco_format["annotations"].append({
                        "id": annotation_id,
                        "image_id": image_id,
                        "category_id": class_id,
                        "bbox": [x_min, y_min, bbox_width, bbox_height],
                        "area": bbox_width * bbox_height,
                        "iscrowd": 0,
                        "segmentation": []  # 不需要分割信息
                    })
                    annotation_id += 1
        
        image_id += 1
    
    # 保存为JSON
    with open(output_json_path, 'w') as f:
        json.dump(coco_format, f)
    print(f"转换完成，COCO格式标注已保存到：{output_json_path}")

# --------------------------
# 修改为你的实际路径
# --------------------------
YOLO_IMAGES_DIR = "E:/APP/文献阅读/DETR/coco128/images"  # 图像文件夹
YOLO_LABELS_DIR = "E:/APP/文献阅读/DETR/coco128/labels"  # 标注文件夹（.txt文件）
OUTPUT_JSON_PATH = "E:/APP/文献阅读/DETR/coco128/annotations/instances_train2017.json"  # 输出JSON路径

# 创建annotations文件夹（如果不存在）
os.makedirs(os.path.dirname(OUTPUT_JSON_PATH), exist_ok=True)

# 执行转换
yolo_to_coco(YOLO_IMAGES_DIR, YOLO_LABELS_DIR, OUTPUT_JSON_PATH)
